{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# default_exp bert_preprocessing.create_bert_features\n",
    "import os\n",
    "%load_ext autoreload\n",
    "%autoreload 2\n",
    "os.environ[\"CUDA_VISIBLE_DEVICES\"] = \"-1\"\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Create Bert Features\n",
    "\n",
    "Module to convert preprocessing function outputs to model input, including\n",
    "\n",
    "- convert texts to token id using transformers tokenizers\n",
    "- convert labels to label ids / token ids"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Imports"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# export\n",
    "import random\n",
    "\n",
    "import numpy as np\n",
    "from typing import Callable, Dict, Union, Tuple\n",
    "\n",
    "from loguru import logger\n",
    "\n",
    "from m3tl.special_tokens import PREDICT\n",
    "from m3tl.bert_preprocessing.bert_utils import (create_instances_from_document)\n",
    "from m3tl.base_params import BaseParams\n",
    "\n",
    "from transformers import PreTrainedTokenizer\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# export\n",
    "\n",
    "\n",
    "def convert_labels_to_ids(\n",
    "        target: dict,\n",
    "        problem: str,\n",
    "        problem_type: str,\n",
    "        label_encoder,\n",
    "        tokenizer=None,\n",
    "        decoding_length: int = None,\n",
    "        custom_label_handling_fn: Callable = None,\n",
    "        tokenized_inputs: dict = None,\n",
    "        **kwargs) -> tuple:\n",
    "\n",
    "    if isinstance(target, dict):\n",
    "        if len(target) == 1 and 'labels' in target:\n",
    "            target = target['labels']\n",
    "\n",
    "    if custom_label_handling_fn is not None:\n",
    "        res = custom_label_handling_fn(\n",
    "            target, label_encoder, tokenizer, decoding_length, tokenized_inputs=tokenized_inputs, problem=problem, **kwargs)\n",
    "\n",
    "    # labelid, label_mask tuple\n",
    "    if isinstance(res, tuple):\n",
    "        if res[1] is None:\n",
    "            return {'{}_label_ids'.format(problem): res[0]}\n",
    "        return {'{}_mask'.format(problem): res[1], '{}_label_ids'.format(problem): res[0]}\n",
    "    elif isinstance(res, dict):\n",
    "        return res\n",
    "\n",
    "    return {'{}_label_ids'.format(problem): res}\n",
    "\n",
    "\n",
    "def create_bert_pretraining(problem,\n",
    "                            inputs_list,\n",
    "                            label_encoder,\n",
    "                            params,\n",
    "                            tokenizer\n",
    "                            ):\n",
    "    \"\"\"Slight modification of original code\n",
    "\n",
    "    Raises:\n",
    "        ValueError -- Input format not right\n",
    "    \"\"\"\n",
    "    if not isinstance(inputs_list[0][0], list):\n",
    "        raise ValueError('inputs is expected to be list of list of list.')\n",
    "\n",
    "    all_documents = []\n",
    "    for document in inputs_list:\n",
    "        all_documents.append([])\n",
    "        doc, _ = document\n",
    "        for sentence in doc:\n",
    "            all_documents[-1].append(tokenizer.tokenize('\\t'.join(sentence)))\n",
    "\n",
    "    all_documents = [d for d in all_documents if d]\n",
    "    rng = random.Random()\n",
    "    rng.shuffle(all_documents)\n",
    "\n",
    "    vocab_words = list(tokenizer.vocab.keys())\n",
    "    instances = []\n",
    "\n",
    "    print_count = 0\n",
    "    return_list = []\n",
    "    for _ in range(params.dupe_factor):\n",
    "        for document_index in range(len(all_documents)):\n",
    "            instances = create_instances_from_document(\n",
    "                all_documents,\n",
    "                document_index,\n",
    "                params.max_seq_len,\n",
    "                params.short_seq_prob,\n",
    "                params.masked_lm_prob,\n",
    "                params.max_predictions_per_seq,\n",
    "                vocab_words, rng)\n",
    "            for instance in instances:\n",
    "                tokens = instance.tokens\n",
    "                segment_ids = list(instance.segment_ids)\n",
    "\n",
    "                masked_lm_positions = list(instance.masked_lm_positions)\n",
    "\n",
    "                next_sentence_label = 1 if instance.is_random_next else 0\n",
    "\n",
    "                mask_lm_dict = tokenizer(instance.masked_lm_labels,\n",
    "                                         truncation=False,\n",
    "                                         is_split_into_words=True,\n",
    "                                         padding='max_length',\n",
    "                                         max_length=params.max_predictions_per_seq,\n",
    "                                         return_special_tokens_mask=False,\n",
    "                                         add_special_tokens=False)\n",
    "                input_ids = tokenizer.convert_tokens_to_ids(tokens)\n",
    "                input_mask = [1 for _ in input_ids]\n",
    "                masked_lm_ids = mask_lm_dict['input_ids']\n",
    "                masked_lm_weights = mask_lm_dict['attention_mask']\n",
    "                masked_lm_positions = masked_lm_positions + \\\n",
    "                    masked_lm_ids[len(masked_lm_positions):]\n",
    "\n",
    "                assert len(input_ids) == len(\n",
    "                    segment_ids), (len(input_ids), len(segment_ids))\n",
    "                assert len(masked_lm_ids) == len(masked_lm_positions), (len(\n",
    "                    masked_lm_ids), len(masked_lm_positions))\n",
    "\n",
    "                yield_dict = {\n",
    "                    \"input_ids\": input_ids,\n",
    "                    \"input_mask\": input_mask,\n",
    "                    \"segment_ids\": segment_ids,\n",
    "                    \"masked_lm_positions\": masked_lm_positions,\n",
    "                    \"masked_lm_ids\": masked_lm_ids,\n",
    "                    \"masked_lm_weights\": masked_lm_weights,\n",
    "                    \"next_sentence_label_ids\": next_sentence_label\n",
    "                }\n",
    "\n",
    "                if print_count < 3:\n",
    "                    logger.debug('%s : %s' %\n",
    "                                               ('tokens', ' '.join([str(x) for x in tokens])))\n",
    "                    for k, v in yield_dict.items():\n",
    "                        if not isinstance(v, int):\n",
    "                            logger.debug('%s : %s' %\n",
    "                                                       (k, ' '.join([str(x) for x in v])))\n",
    "                    print_count += 1\n",
    "\n",
    "                return_list.append(yield_dict)\n",
    "    return return_list\n",
    "\n",
    "\n",
    "def mask_inputs_for_mask_lm(inp_text: str, tokenizer: PreTrainedTokenizer, mask_prob=0.1, max_length=128, max_predictions_per_seq=20) -> str:\n",
    "    if not inp_text:\n",
    "        return None, None\n",
    "    inp_text = list(inp_text)\n",
    "    mask_idx = [i for i in range(min(len(inp_text), max_length))\n",
    "                if random.uniform(0, 1) <= mask_prob]\n",
    "    if not mask_idx:\n",
    "        return None, None\n",
    "    masked_text = [inp_text[i] for i in mask_idx]\n",
    "    inp_text = [t if i not in mask_idx else '[MASK]' for i,\n",
    "                t in enumerate(inp_text)]\n",
    "\n",
    "    tokenized_dict = tokenizer(\n",
    "        inp_text, None,\n",
    "        truncation=True,\n",
    "        max_length=max_length,\n",
    "        is_split_into_words=True,\n",
    "        padding=False,\n",
    "        return_special_tokens_mask=False,\n",
    "        add_special_tokens=True,\n",
    "        return_overflowing_tokens=True)\n",
    "\n",
    "    # create mask lm features\n",
    "    mask_lm_dict = tokenizer(masked_text,\n",
    "                             truncation=True,\n",
    "                             is_split_into_words=True,\n",
    "                             padding='max_length',\n",
    "                             max_length=max_predictions_per_seq,\n",
    "                             return_special_tokens_mask=False,\n",
    "                             add_special_tokens=False,)\n",
    "\n",
    "    mask_token_id = tokenizer(\n",
    "        '[MASK]', add_special_tokens=False, is_split_into_words=False)['input_ids'][0]\n",
    "    masked_lm_positions = [i for i, input_id in enumerate(\n",
    "        tokenized_dict['input_ids']) if input_id == mask_token_id]\n",
    "    # pad masked_lm_positions to max_predictions_per_seq\n",
    "    if len(masked_lm_positions) < max_predictions_per_seq:\n",
    "        masked_lm_positions = masked_lm_positions + \\\n",
    "            [0 for _ in range(max_predictions_per_seq -\n",
    "                              len(masked_lm_positions))]\n",
    "    masked_lm_positions = masked_lm_positions[:max_predictions_per_seq]\n",
    "    masked_lm_ids = np.array(mask_lm_dict['input_ids'], dtype='int32')\n",
    "    masked_lm_weights = np.array(mask_lm_dict['attention_mask'], dtype='int32')\n",
    "    mask_lm_dict = {'masked_lm_positions': masked_lm_positions,\n",
    "                    'masked_lm_ids': masked_lm_ids,\n",
    "                    'masked_lm_weights': masked_lm_weights}\n",
    "\n",
    "    return tokenized_dict, mask_lm_dict\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Shortcut functions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# export\n",
    "\n",
    "def collect_modal_name(inputs: dict) -> list:\n",
    "    return [k for k in inputs.keys() if 'modal_type' not in k and 'modal_info' not in k]\n",
    "\n",
    "\n",
    "def infer_modal(inputs: dict) -> dict:\n",
    "    modal_name_list = collect_modal_name(inputs)\n",
    "    for modal_name in modal_name_list:\n",
    "        modal_type_name = '{}_modal_type'.format(modal_name)\n",
    "        # infer modal type\n",
    "        if modal_type_name not in inputs:\n",
    "            inp = inputs[modal_name]\n",
    "\n",
    "            if isinstance(inp, str):\n",
    "                inputs[modal_type_name] = 'text'\n",
    "            elif isinstance(inp, int):\n",
    "                inputs[modal_type_name] = 'category'\n",
    "            elif isinstance(inp, list) and len(inp) == 0:\n",
    "                raise ValueError(\n",
    "                    'Found empty list in input but modal_type is not provided: {}'.format(inputs))\n",
    "            # tokenized text\n",
    "            elif isinstance(inp, list) and isinstance(inp[0], str):\n",
    "                inputs[modal_type_name] = 'text'\n",
    "            elif isinstance(inp, list) and isinstance(inp[0], int):\n",
    "                inputs[modal_type_name] = 'category'\n",
    "            else:\n",
    "                inputs[modal_type_name] = 'array'\n",
    "    return inputs\n",
    "\n",
    "\n",
    "def unify_inputs(inputs: Union[Dict, str]) -> Dict:\n",
    "    \"\"\"\n",
    "    unify various input format to single one as following:\n",
    "    {\n",
    "        'title': 'this is test',\n",
    "        'image': np.array([1,2,3]),\n",
    "        'class': 'a',\n",
    "        'title_modal_type': 'text', \n",
    "        'image_modal_type': 'array',\n",
    "        'class_modal_type': 'category',\n",
    "        'class_modal_info': 10\n",
    "    }\n",
    "    \"\"\"\n",
    "\n",
    "    # inputs is single-modal\n",
    "    if isinstance(inputs, str) or isinstance(inputs, list) or 'a' in inputs:\n",
    "        return {\n",
    "            'text': inputs,\n",
    "            'text_modal_type': 'text'\n",
    "        }\n",
    "    inputs = infer_modal(inputs)\n",
    "    return inputs\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'text': 'this is a test',\n",
       " 'image': array([1, 2, 3]),\n",
       " 'class': 0,\n",
       " 'text_modal_type': 'text',\n",
       " 'image_modal_type': 'array',\n",
       " 'class_modal_type': 'category'}"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# test pure text input\n",
    "unify_inputs('this is a test')\n",
    "\n",
    "# test legacy multimodal input\n",
    "unify_inputs({'text': 'this is a test',\n",
    "             'image': np.array([1, 2, 3]), 'class': 0})\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Modal handling functions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# export\n",
    "class NextExampleException(Exception):\n",
    "    pass\n",
    "\n",
    "\n",
    "def text_modal_input_handling(\n",
    "        problem: str,\n",
    "        modal_name: str,\n",
    "        modal_inputs: dict,\n",
    "        target: dict = None,\n",
    "        params: BaseParams = None,\n",
    "        tokenizer: PreTrainedTokenizer = None) -> Tuple[dict, dict]:\n",
    "    is_mask_lm = (\n",
    "        params.problem_type[problem] == 'masklm'\n",
    "        if problem in params.problem_type else False\n",
    "    )\n",
    "    is_seq = (\n",
    "        params.problem_type[problem] == 'seq_tag'\n",
    "        if problem in params.problem_type else False\n",
    "    )\n",
    "    # tokenize inputs, now the length is fixed, target == raw_target\n",
    "    if isinstance(modal_inputs, dict):\n",
    "        tokens_a = modal_inputs['a']\n",
    "        tokens_b = modal_inputs['b']\n",
    "    else:\n",
    "        tokens_a = modal_inputs\n",
    "        tokens_b = None\n",
    "\n",
    "    if is_mask_lm:\n",
    "        tokenized_dict, mlm_feature_dict = mask_inputs_for_mask_lm(\n",
    "            tokens_a, tokenizer, mask_prob=params.masked_lm_prob,\n",
    "            max_length=params.max_seq_len, max_predictions_per_seq=params.max_predictions_per_seq)\n",
    "        if tokenized_dict is None:\n",
    "            # hacky approach to continue outer loop\n",
    "            raise NextExampleException\n",
    "    else:\n",
    "        mlm_feature_dict = {}\n",
    "\n",
    "        tokenized_dict = tokenizer.encode_plus(\n",
    "            tokens_a, tokens_b,\n",
    "            truncation=True,\n",
    "            max_length=params.max_seq_len,\n",
    "            is_split_into_words=False,\n",
    "            padding=False,\n",
    "            return_special_tokens_mask=is_seq,\n",
    "            add_special_tokens=True,\n",
    "            return_overflowing_tokens=True,\n",
    "            return_token_type_ids=True)\n",
    "\n",
    "    input_ids = tokenized_dict['input_ids']\n",
    "    segment_ids = tokenized_dict['token_type_ids']\n",
    "    input_mask = tokenized_dict['attention_mask']\n",
    "\n",
    "    modal_feature_dict = {\n",
    "        '{}_input_ids'.format(modal_name): input_ids,\n",
    "        '{}_mask'.format(modal_name): input_mask,\n",
    "        '{}_segment_ids'.format(modal_name): segment_ids\n",
    "    }\n",
    "\n",
    "    modal_feature_dict.update(mlm_feature_dict)\n",
    "    return modal_feature_dict, tokenized_dict\n",
    "\n",
    "\n",
    "def array_modal_input_handling(\n",
    "        problem: str,\n",
    "        modal_name: str,\n",
    "        modal_inputs: dict,\n",
    "        target: dict = None,\n",
    "        params: BaseParams = None,\n",
    "        tokenizer: PreTrainedTokenizer = None) -> dict:\n",
    "    modal_inputs = np.array(modal_inputs)\n",
    "    # avoid unbounded var\n",
    "    tokenized_dict = None\n",
    "\n",
    "    if len(modal_inputs.shape) == 1:\n",
    "        modal_inputs = np.expand_dims(modal_inputs, axis=0)\n",
    "    # maybe we need to add a,b support for multimodal inputs?\n",
    "    segment_ids = np.zeros(\n",
    "        modal_inputs.shape[0], dtype=np.int32)\n",
    "    input_mask = [1]*len(modal_inputs)\n",
    "    modal_feature_dict = {\n",
    "        '{}_input_ids'.format(modal_name): modal_inputs,\n",
    "        '{}_mask'.format(modal_name): input_mask,\n",
    "        '{}_segment_ids'.format(modal_name): segment_ids}\n",
    "    return modal_feature_dict\n",
    "\n",
    "\n",
    "def category_modal_input_handling(\n",
    "        problem: str,\n",
    "        modal_name: str,\n",
    "        modal_inputs: dict,\n",
    "        target: dict = None,\n",
    "        params: BaseParams = None,\n",
    "        tokenizer: PreTrainedTokenizer = None,\n",
    "        raw_inputs: dict = None) -> dict:\n",
    "\n",
    "    modal_inputs = np.array(modal_inputs, dtype='int')\n",
    "\n",
    "    modal_info_name = '{}_modal_info'.format(modal_name)\n",
    "    if modal_info_name not in raw_inputs:\n",
    "        raise KeyError(\n",
    "            \"category modal {} modal info not set\".format(modal_name))\n",
    "    params.set_problem_info(\n",
    "        problem=problem, info_name=modal_info_name, info=raw_inputs[modal_info_name])\n",
    "\n",
    "    if len(modal_inputs.shape) == 0:\n",
    "        modal_inputs = np.expand_dims(modal_inputs, axis=0)\n",
    "    elif len(modal_inputs) > 2:\n",
    "        raise ValueError(\n",
    "            'rank of category modal should be less than 2, get:{} of problem {}'.format(\n",
    "                len(modal_inputs.shape), problem))\n",
    "\n",
    "    segment_ids = np.zeros(\n",
    "        modal_inputs.shape[0], dtype=np.int32)\n",
    "    input_mask = [1]*len(modal_inputs)\n",
    "\n",
    "    modal_feature_dict = {\n",
    "        '{}_input_ids'.format(modal_name): modal_inputs,\n",
    "        '{}_mask'.format(modal_name): input_mask,\n",
    "        '{}_segment_ids'.format(modal_name): segment_ids}\n",
    "    return modal_feature_dict\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "WARNING:root:bert_config not exists. will load model from huggingface checkpoint.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Adding new problem weibo_fake_ner, problem type: seq_tag\n",
      "Adding new problem weibo_fake_multi_cls, problem type: multi_cls\n",
      "Adding new problem weibo_fake_cls, problem type: cls\n",
      "Adding new problem weibo_masklm, problem type: masklm\n",
      "Adding new problem weibo_fake_regression, problem type: regression\n",
      "Adding new problem weibo_fake_vector_fit, problem type: vector_fit\n",
      "Adding new problem weibo_premask_mlm, problem type: premask_mlm\n"
     ]
    }
   ],
   "source": [
    "# test seq_tag\n",
    "from m3tl.test_base import TestBase\n",
    "from m3tl.utils import load_transformer_tokenizer\n",
    "from m3tl.predefined_problems import generate_fake_data\n",
    "tb = TestBase()\n",
    "params = tb.params\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "({'desc_input_ids': [101, 8554, 8310, 143, 10060, 102], 'desc_mask': [1, 1, 1, 1, 1, 1], 'desc_segment_ids': [0, 0, 0, 0, 0, 0]}, {'overflowing_tokens': [], 'num_truncated_tokens': -250, 'input_ids': [101, 8554, 8310, 143, 10060, 102], 'token_type_ids': [0, 0, 0, 0, 0, 0], 'special_tokens_mask': [1, 0, 0, 0, 0, 1], 'attention_mask': [1, 1, 1, 1, 1, 1]})\n",
      "{'inputs_cate_input_ids': array([0]), 'inputs_cate_mask': [1], 'inputs_cate_segment_ids': array([0], dtype=int32)}\n"
     ]
    }
   ],
   "source": [
    "problem = 'weibo_fake_ner'\n",
    "modal_name = 'desc'\n",
    "modal_inputs = 'this is a test'\n",
    "target = None\n",
    "tokenizer = load_transformer_tokenizer(\n",
    "    tokenizer_name=params.transformer_tokenizer_name, load_module_name=params.transformer_tokenizer_loading)\n",
    "print(text_modal_input_handling(\n",
    "    problem=problem, modal_name=modal_name, modal_inputs=modal_inputs, tokenizer=tokenizer, params=params))\n",
    "fake_data = generate_fake_data('dict')\n",
    "print(category_modal_input_handling(\n",
    "    problem=problem, modal_name='inputs_cate', modal_inputs=fake_data[0]['inputs_cate'], tokenizer=tokenizer, params=params, raw_inputs=fake_data[0]))\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Create bert multimodal features"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "# export\n",
    "\n",
    "def separate_inputs_labels(inp: dict) -> Tuple[dict, dict]:\n",
    "    inputs_dict = {k: v for k, v in inp.items() if k.startswith('inputs')}\n",
    "    labels_dict = {k: v for k, v in inp.items() if k.startswith('labels')}\n",
    "\n",
    "    inputs_dict = remove_dict_prefix(inputs_dict, key='inputs_')\n",
    "    labels_dict = remove_dict_prefix(labels_dict, key='labels_')\n",
    "    # some special handling for labels here\n",
    "    if not labels_dict:\n",
    "        labels_dict = None\n",
    "\n",
    "    elif len(labels_dict) == 1:\n",
    "        labels_dict = list(labels_dict.values())[0]\n",
    "    return inputs_dict, labels_dict\n",
    "\n",
    "\n",
    "def remove_dict_prefix(inp: dict, key='inputs_') -> dict:\n",
    "\n",
    "    cleaned_cols = [\n",
    "        c.replace(key, '') for c in inp.keys()]\n",
    "\n",
    "    return {\n",
    "        cleaned_column: inp[in_column]\n",
    "        for cleaned_column, in_column in zip(\n",
    "            cleaned_cols, inp.keys()\n",
    "        )\n",
    "    }\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# export\n",
    "\n",
    "def _create_multimodal_bert_features(problem,\n",
    "                                     example_list,\n",
    "                                     label_encoder,\n",
    "                                     params,\n",
    "                                     tokenizer,\n",
    "                                     mode,\n",
    "                                     problem_type,\n",
    "                                     is_seq):\n",
    "    if problem_type == 'pretrain':\n",
    "        return create_bert_pretraining(\n",
    "            problem=problem,\n",
    "            inputs_list=example_list,\n",
    "            label_encoder=label_encoder,\n",
    "            params=params,\n",
    "            tokenizer=tokenizer\n",
    "        )\n",
    "\n",
    "    is_mask_lm = problem_type == 'masklm'\n",
    "\n",
    "    for example_id, example in enumerate(example_list):\n",
    "        raw_inputs, raw_target = separate_inputs_labels(example)\n",
    "\n",
    "        raw_inputs = unify_inputs(raw_inputs)\n",
    "        modal_name_list = [m for m in collect_modal_name(\n",
    "            raw_inputs) if m != 'record_id']\n",
    "\n",
    "        # save problem modal type\n",
    "        if example_id == 0:\n",
    "            for modal_name in modal_name_list:\n",
    "                params.set_problem_info(problem=problem, info_name='{}_modal_type'.format(\n",
    "                    modal_name), info=raw_inputs['{}_modal_type'.format(modal_name)])\n",
    "                modal_info_name = '{}_modal_info'.format(modal_name)\n",
    "                if modal_info_name in raw_inputs:\n",
    "                    params.set_problem_info(\n",
    "                        problem=problem, info_name=modal_info_name, info=raw_inputs[modal_info_name])\n",
    "\n",
    "        return_dict = {}\n",
    "        if 'record_id' in raw_inputs:\n",
    "            return_dict['record_id'] = raw_inputs['record_id']\n",
    "        try:\n",
    "            for idx, modal_name in enumerate(modal_name_list):\n",
    "\n",
    "                modal_inputs = raw_inputs[modal_name]\n",
    "                modal_type = raw_inputs['{}_modal_type'.format(modal_name)]\n",
    "\n",
    "                target = raw_target\n",
    "                if modal_type == 'text':\n",
    "                    modal_feature_dict, tokenized_dict = text_modal_input_handling(\n",
    "                        problem=problem,\n",
    "                        modal_name=modal_name,\n",
    "                        modal_inputs=modal_inputs,\n",
    "                        target=target,\n",
    "                        params=params,\n",
    "                        tokenizer=tokenizer\n",
    "                    )\n",
    "\n",
    "                elif modal_type == 'array':\n",
    "                    modal_feature_dict = array_modal_input_handling(\n",
    "                        problem=problem,\n",
    "                        modal_name=modal_name,\n",
    "                        modal_inputs=modal_inputs,\n",
    "                        target=target,\n",
    "                        params=params,\n",
    "                        tokenizer=tokenizer)\n",
    "                    tokenized_dict = None\n",
    "                elif modal_type == 'category':\n",
    "                    modal_feature_dict = category_modal_input_handling(\n",
    "                        problem=problem,\n",
    "                        modal_name=modal_name,\n",
    "                        modal_inputs=modal_inputs,\n",
    "                        target=target,\n",
    "                        params=params,\n",
    "                        tokenizer=tokenizer,\n",
    "                        raw_inputs=raw_inputs\n",
    "                    )\n",
    "                    tokenized_dict = None\n",
    "                else:\n",
    "                    raise ValueError(\n",
    "                        'Undefined modal type: {}'.format(modal_type))\n",
    "\n",
    "                # encode labels\n",
    "                if mode != PREDICT:\n",
    "                    if not is_mask_lm:\n",
    "                        custom_label_handling_fn = params.label_handling_fn.get(\n",
    "                            problem_type, None)\n",
    "                        label_dict = convert_labels_to_ids(\n",
    "                            target=target, problem=problem, problem_type=problem_type,\n",
    "                            label_encoder=label_encoder, tokenizer=tokenizer,\n",
    "                            decoding_length=params.decode_max_seq_len,\n",
    "                            custom_label_handling_fn=custom_label_handling_fn,\n",
    "                            tokenized_inputs=tokenized_dict,\n",
    "                            modal_name=modal_name,\n",
    "                            modal_type=modal_type)\n",
    "\n",
    "                        modal_feature_dict.update(label_dict)\n",
    "                return_dict.update(modal_feature_dict)\n",
    "\n",
    "        except NextExampleException:\n",
    "            continue\n",
    "\n",
    "        if params.print_data:\n",
    "            if example_id < 1:\n",
    "                if isinstance(raw_inputs, dict):\n",
    "                    for raw_input_name, raw_input in raw_inputs.items():\n",
    "                        logger.debug('{}: {}'.format(\n",
    "                            raw_input_name, str(raw_input)[:200]))\n",
    "                else:\n",
    "                    logger.debug(str(raw_inputs)[:200])\n",
    "                for return_key, return_item in return_dict.items():\n",
    "                    logger.debug('{}: {}'.format(\n",
    "                        return_key, str(return_item)[:200]))\n",
    "        yield return_dict\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Wrapper functions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# export\n",
    "def create_multimodal_bert_features(problem,\n",
    "                                    example_list,\n",
    "                                    label_encoder,\n",
    "                                    params,\n",
    "                                    tokenizer,\n",
    "                                    mode,\n",
    "                                    problem_type,\n",
    "                                    is_seq):\n",
    "    if problem_type == 'pretrain':\n",
    "        return create_bert_pretraining(\n",
    "            problem=problem,\n",
    "            inputs_list=example_list,\n",
    "            label_encoder=label_encoder,\n",
    "            params=params,\n",
    "            tokenizer=tokenizer\n",
    "        )\n",
    "    gen = _create_multimodal_bert_features(problem,\n",
    "                                           example_list,\n",
    "                                           label_encoder,\n",
    "                                           params,\n",
    "                                           tokenizer,\n",
    "                                           mode,\n",
    "                                           problem_type,\n",
    "                                           is_seq)\n",
    "    return_dict_list = [d for d in gen]\n",
    "    return return_dict_list\n",
    "\n",
    "\n",
    "def create_multimodal_bert_features_generator(problem,\n",
    "                                              example_list,\n",
    "                                              label_encoder,\n",
    "                                              params,\n",
    "                                              tokenizer,\n",
    "                                              mode,\n",
    "                                              problem_type,\n",
    "                                              is_seq):\n",
    "    if problem_type == 'pretrain':\n",
    "        return create_bert_pretraining(\n",
    "            problem=problem,\n",
    "            inputs_list=example_list,\n",
    "            label_encoder=label_encoder,\n",
    "            params=params,\n",
    "            tokenizer=tokenizer\n",
    "        )\n",
    "    gen = _create_multimodal_bert_features(problem,\n",
    "                                           example_list,\n",
    "                                           label_encoder,\n",
    "                                           params,\n",
    "                                           tokenizer,\n",
    "                                           mode,\n",
    "                                           problem_type,\n",
    "                                           is_seq)\n",
    "    return gen\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {},
 "nbformat": 4,
 "nbformat_minor": 2
}
